import os
import random
import shutil

# 原始文件夹路径
image_folder = r"D:\songlin\data\8.图像分割相关\马\image"
yolo_folder = r"D:\songlin\data\8.图像分割相关\马\yolo_txt"

# 创建目标文件夹结构
dataset_folder = r"D:\songlin\data\8.图像分割相关\马\dataset"
train_image_folder = os.path.join(dataset_folder, "train", "images")
train_label_folder = os.path.join(dataset_folder, "train", "labels")
val_image_folder = os.path.join(dataset_folder, "val", "images")
val_label_folder = os.path.join(dataset_folder, "val", "labels")

# 创建所有必要的文件夹
os.makedirs(train_image_folder, exist_ok=True)
os.makedirs(train_label_folder, exist_ok=True)
os.makedirs(val_image_folder, exist_ok=True)
os.makedirs(val_label_folder, exist_ok=True)

# 获取所有图片文件名（不带扩展名）
image_files = [os.path.splitext(f)[0] for f in os.listdir(image_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]

# 随机打乱文件列表
random.shuffle(image_files)

# 计算分割点
split_idx = int(len(image_files) * 0.8)
train_files = image_files[:split_idx]
val_files = image_files[split_idx:]

# 复制训练集文件
for filename in train_files:
    # 复制图片文件
    for ext in ['.jpg', '.jpeg', '.png']:
        src_img = os.path.join(image_folder, filename + ext)
        if os.path.exists(src_img):
            shutil.copy(src_img, train_image_folder)
            break

    # 复制标注文件
    src_txt = os.path.join(yolo_folder, filename + ".txt")
    if os.path.exists(src_txt):
        shutil.copy(src_txt, train_label_folder)

# 复制验证集文件
for filename in val_files:
    # 复制图片文件
    for ext in ['.jpg', '.jpeg', '.png']:
        src_img = os.path.join(image_folder, filename + ext)
        if os.path.exists(src_img):
            shutil.copy(src_img, val_image_folder)
            break

    # 复制标注文件
    src_txt = os.path.join(yolo_folder, filename + ".txt")
    if os.path.exists(src_txt):
        shutil.copy(src_txt, val_label_folder)

print(f"数据集划分完成！训练集: {len(train_files)} 张, 验证集: {len(val_files)} 张")